"""
自回归移动平均模型
"""
import numpy as np
import pandas as pd
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import TimeSeriesSplit
from statsmodels.tsa.api import statespace


def arima(series, order=(2, 1, 4), seasonal_order=(0, 1, 1, 7)):
    """
    自回归移动平均模型
    :param series:
    :param order:
    :param seasonal_order:
    :return:
    """
    if isinstance(series, pd.Series):
        data_set = series.to_numpy()

    if not isinstance(series, np.ndarray):
        data_set = np.array(series)
    else:
        data_set = series

    rms_arima = []

    tscv = TimeSeriesSplit()
    for train_idx, test_idx in tscv.split(data_set):
        train, test = data_set[train_idx], data_set[test_idx]

        # arima
        """
        order：
          AR参数，差异和MA参数的数量的模型的（p，d，q）顺序。
          默认值为AR（1）模型：（1,0,0）
        seasonal_order：
          AR参数，差异，MA参数和周期性的模型季节性成分的（P，D，Q，s）顺序。
          s是给出周期性（季节的周期数）的整数，对于季度数据而言通常为4，
          对于月度数据而言通常为12。 默认为无季节性影响
        """
        arima_model = statespace.SARIMAX(data_set,
                                         order=order,
                                         seasonal_order=seasonal_order).fit()
        y_pred = arima_model.predict(start=test_idx[0],
                                     end=test_idx[-1],
                                     dynamic=True)
        rms_arima.append(np.sqrt(mean_squared_error(test, y_pred)))

    print('rms: {}'.format(np.mean(rms_arima)))
